{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jupyter/.local/lib/python3.5/site-packages/ipykernel_launcher.py:79: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n"
     ]
    }
   ],
   "source": [
    "import os, time\n",
    "import matplotlib.pyplot as plt\n",
    "import itertools\n",
    "import pickle\n",
    "import imageio\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import torch.optim as optim\n",
    "from torchvision import datasets, transforms\n",
    "from torch.autograd import Variable\n",
    "from torch.utils.data.dataloader import default_collate\n",
    "\n",
    "# G(z)\n",
    "class generator(nn.Module):\n",
    "    # initializers\n",
    "    def __init__(self, d=128):\n",
    "        super(generator, self).__init__()\n",
    "        self.deconv1 = nn.ConvTranspose2d(100, d*8, 4, 1, 0)\n",
    "        self.deconv1_bn = nn.BatchNorm2d(d*8)\n",
    "        self.deconv2 = nn.ConvTranspose2d(d*8, d*4, 4, 2, 1)\n",
    "        self.deconv2_bn = nn.BatchNorm2d(d*4)\n",
    "        self.deconv3 = nn.ConvTranspose2d(d*4, d*2, 4, 2, 1)\n",
    "        self.deconv3_bn = nn.BatchNorm2d(d*2)\n",
    "        self.deconv4 = nn.ConvTranspose2d(d*2, d, 4, 2, 1)\n",
    "        self.deconv4_bn = nn.BatchNorm2d(d)\n",
    "        self.deconv5 = nn.ConvTranspose2d(d, 1, 4, 2, 1)\n",
    "\n",
    "    # weight_init\n",
    "    def weight_init(self, mean, std):\n",
    "        for m in self._modules:\n",
    "            normal_init(self._modules[m], mean, std)\n",
    "\n",
    "    # forward method\n",
    "    def forward(self, input):\n",
    "        # x = F.relu(self.deconv1(input))\n",
    "        x = F.relu(self.deconv1_bn(self.deconv1(input)))\n",
    "        x = F.relu(self.deconv2_bn(self.deconv2(x)))\n",
    "        x = F.relu(self.deconv3_bn(self.deconv3(x)))\n",
    "        x = F.relu(self.deconv4_bn(self.deconv4(x)))\n",
    "        x = F.tanh(self.deconv5(x))\n",
    "\n",
    "        return x\n",
    "\n",
    "class discriminator(nn.Module):\n",
    "    # initializers\n",
    "    def __init__(self, d=128):\n",
    "        super(discriminator, self).__init__()\n",
    "        self.conv1 = nn.Conv2d(1, d, 4, 2, 1)\n",
    "        self.conv2 = nn.Conv2d(d, d*2, 4, 2, 1)\n",
    "        self.conv2_bn = nn.BatchNorm2d(d*2)\n",
    "        self.conv3 = nn.Conv2d(d*2, d*4, 4, 2, 1)\n",
    "        self.conv3_bn = nn.BatchNorm2d(d*4)\n",
    "        self.conv4 = nn.Conv2d(d*4, d*8, 4, 2, 1)\n",
    "        self.conv4_bn = nn.BatchNorm2d(d*8)\n",
    "        self.conv5 = nn.Conv2d(d*8, 1, 4, 1, 0)\n",
    "\n",
    "    # weight_init\n",
    "    def weight_init(self, mean, std):\n",
    "        for m in self._modules:\n",
    "            normal_init(self._modules[m], mean, std)\n",
    "\n",
    "    # forward method\n",
    "    def forward(self, input):\n",
    "        x = F.leaky_relu(self.conv1(input), 0.2)\n",
    "        x = F.leaky_relu(self.conv2_bn(self.conv2(x)), 0.2)\n",
    "        x = F.leaky_relu(self.conv3_bn(self.conv3(x)), 0.2)\n",
    "        x = F.leaky_relu(self.conv4_bn(self.conv4(x)), 0.2)\n",
    "        x = F.sigmoid(self.conv5(x))\n",
    "\n",
    "        return x\n",
    "\n",
    "def normal_init(m, mean, std):\n",
    "    if isinstance(m, nn.ConvTranspose2d) or isinstance(m, nn.Conv2d):\n",
    "        m.weight.data.normal_(mean, std)\n",
    "        m.bias.data.zero_()\n",
    "\n",
    "fixed_z_ = torch.randn((5 * 5, 100)).view(-1, 100, 1, 1)    # fixed noise\n",
    "fixed_z_ = Variable(fixed_z_.cuda(), volatile=True)\n",
    "\n",
    "def show_result(num_epoch, show = False, save = False, path = 'result.png', isFix=False):\n",
    "    z_ = torch.randn((5*5, 100)).view(-1, 100, 1, 1)\n",
    "    z_ = Variable(z_.cuda(), volatile=True)\n",
    "\n",
    "    G.eval()\n",
    "    if isFix:\n",
    "        test_images = G(fixed_z_)\n",
    "    else:\n",
    "        test_images = G(z_)\n",
    "    G.train()\n",
    "\n",
    "    size_figure_grid = 5\n",
    "    fig, ax = plt.subplots(size_figure_grid, size_figure_grid, figsize=(5, 5))\n",
    "    for i, j in itertools.product(range(size_figure_grid), range(size_figure_grid)):\n",
    "        ax[i, j].get_xaxis().set_visible(False)\n",
    "        ax[i, j].get_yaxis().set_visible(False)\n",
    "\n",
    "    for k in range(5*5):\n",
    "        i = k // 5\n",
    "        j = k % 5\n",
    "        ax[i, j].cla()\n",
    "        ax[i, j].imshow(test_images[k, 0].cpu().data.numpy(), cmap='gray')\n",
    "\n",
    "    label = 'Epoch {0}'.format(num_epoch)\n",
    "    fig.text(0.5, 0.04, label, ha='center')\n",
    "    plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()\n",
    "\n",
    "def show_train_hist(hist, show = False, save = False, path = 'Train_hist.png'):\n",
    "    x = range(len(hist['D_losses']))\n",
    "\n",
    "    y1 = hist['D_losses']\n",
    "    y2 = hist['G_losses']\n",
    "\n",
    "    plt.plot(x, y1, label='D_loss')\n",
    "    plt.plot(x, y2, label='G_loss')\n",
    "\n",
    "    plt.xlabel('Iter')\n",
    "    plt.ylabel('Loss')\n",
    "\n",
    "    plt.legend(loc=4)\n",
    "    plt.grid(True)\n",
    "    plt.tight_layout()\n",
    "\n",
    "    if save:\n",
    "        plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jupyter/.local/lib/python3.5/site-packages/torchvision/transforms/transforms.py:208: UserWarning: The use of the transforms.Scale transform is deprecated, please use transforms.Resize instead.\n",
      "  \"please use transforms.Resize instead.\")\n"
     ]
    }
   ],
   "source": [
    "# training parameters\n",
    "\n",
    "batch_size = 128\n",
    "lr = 0.0002\n",
    "train_epoch = 20\n",
    "\n",
    "img_size = 64\n",
    "\n",
    "def my_collate(batch):\n",
    "    modified_batch = []\n",
    "    for item in batch:\n",
    "        image, label = item\n",
    "        if label is not 9:\n",
    "            modified_batch.append(item)\n",
    "    return default_collate(modified_batch)\n",
    "\n",
    "transform = transforms.Compose([\n",
    "        transforms.Scale(img_size),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(\n",
    "    datasets.FashionMNIST('data', train=True, download=True, transform=transform),\n",
    "    batch_size=batch_size, shuffle=True, collate_fn = my_collate)\n",
    "\n",
    "# network\n",
    "G = generator(128)\n",
    "D = discriminator(128)\n",
    "\n",
    "G.weight_init(mean=0.0, std=0.02)\n",
    "D.weight_init(mean=0.0, std=0.02)\n",
    "\n",
    "G.cuda()\n",
    "D.cuda()\n",
    "\n",
    "# Binary Cross Entropy loss\n",
    "BCE_loss = nn.BCELoss()\n",
    "\n",
    "# Adam optimizer\n",
    "G_optimizer = optim.Adam(G.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "D_optimizer = optim.Adam(D.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "\n",
    "# results save folder\n",
    "if not os.path.isdir('MNIST_DCGAN_results'):\n",
    "    os.mkdir('MNIST_DCGAN_results')\n",
    "if not os.path.isdir('MNIST_DCGAN_results/Random_results'):\n",
    "    os.mkdir('MNIST_DCGAN_results/Random_results')\n",
    "if not os.path.isdir('MNIST_DCGAN_results/Fixed_results'):\n",
    "    os.mkdir('MNIST_DCGAN_results/Fixed_results')\n",
    "\n",
    "train_hist = {}\n",
    "train_hist['D_losses'] = []\n",
    "train_hist['G_losses'] = []\n",
    "train_hist['per_epoch_ptimes'] = []\n",
    "train_hist['total_ptime'] = []\n",
    "num_iter = 0"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.utils.data.dataloader.DataLoader'>\n",
      "tensor([0, 5, 8, 3, 0, 8, 0, 5, 4, 1, 6, 6, 5, 3, 8, 5, 0, 1, 5, 6, 3, 0, 6, 0,\n",
      "        4, 1, 3, 3, 3, 3, 7, 3, 8, 4, 4, 7, 1, 6, 6, 2, 0, 3, 6, 1, 0, 3, 0, 3,\n",
      "        7, 0, 6, 2, 1, 0, 8, 0, 3, 3, 5, 6, 1, 1, 4, 3, 2, 0, 5, 8, 5, 7, 8, 2,\n",
      "        4, 8, 2, 3, 5, 5, 5, 1, 6, 7, 1, 6, 6, 7, 7, 2, 5, 4, 3, 0, 2, 7, 3, 7,\n",
      "        0, 4, 8, 5, 8, 2, 0, 2, 1, 3, 5, 0, 5, 1, 3, 2, 0, 3, 2])\n"
     ]
    }
   ],
   "source": [
    "print(type(train_loader))\n",
    "\n",
    "count = 0\n",
    "\n",
    "for x_, y in train_loader:\n",
    "    count+= 1\n",
    "    # print(x_)\n",
    "    print(y)\n",
    "    if count == 1:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "training start!\n",
      "epoch number 0\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jupyter/.local/lib/python3.5/site-packages/torch/nn/functional.py:1332: UserWarning: nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\n",
      "  warnings.warn(\"nn.functional.sigmoid is deprecated. Use torch.sigmoid instead.\")\n",
      "/home/jupyter/.local/lib/python3.5/site-packages/torch/nn/functional.py:1320: UserWarning: nn.functional.tanh is deprecated. Use torch.tanh instead.\n",
      "  warnings.warn(\"nn.functional.tanh is deprecated. Use torch.tanh instead.\")\n"
     ]
    }
   ],
   "source": [
    "print('training start!')\n",
    "start_time = time.time()\n",
    "for epoch in range(train_epoch):\n",
    "    print(\"epoch number\", epoch)\n",
    "    D_losses = []\n",
    "    G_losses = []\n",
    "    epoch_start_time = time.time()\n",
    "    for x_, _ in train_loader:\n",
    "        # train discriminator D\n",
    "        D.zero_grad()\n",
    "\n",
    "        mini_batch = x_.size()[0]\n",
    "\n",
    "        y_real_ = torch.ones(mini_batch)\n",
    "        y_fake_ = torch.zeros(mini_batch)\n",
    "\n",
    "        x_, y_real_, y_fake_ = Variable(x_.cuda()), Variable(y_real_.cuda()), Variable(y_fake_.cuda())\n",
    "        D_result = D(x_).squeeze()\n",
    "        D_real_loss = BCE_loss(D_result, y_real_)\n",
    "\n",
    "        z_ = torch.randn((mini_batch, 100)).view(-1, 100, 1, 1)\n",
    "        z_ = Variable(z_.cuda())\n",
    "        G_result = G(z_)\n",
    "\n",
    "        D_result = D(G_result).squeeze()\n",
    "        D_fake_loss = BCE_loss(D_result, y_fake_)\n",
    "        D_fake_score = D_result.data.mean()\n",
    "\n",
    "        D_train_loss = D_real_loss + D_fake_loss\n",
    "\n",
    "        D_train_loss.backward()\n",
    "        D_optimizer.step()\n",
    "\n",
    "        # D_losses.append(D_train_loss.data[0])\n",
    "        D_losses.append(D_train_loss.data)\n",
    "\n",
    "        # train generator G\n",
    "        G.zero_grad()\n",
    "\n",
    "        z_ = torch.randn((mini_batch, 100)).view(-1, 100, 1, 1)\n",
    "        z_ = Variable(z_.cuda())\n",
    "\n",
    "        G_result = G(z_)\n",
    "        D_result = D(G_result).squeeze()\n",
    "        G_train_loss = BCE_loss(D_result, y_real_)\n",
    "        G_train_loss.backward()\n",
    "        G_optimizer.step()\n",
    "\n",
    "        G_losses.append(G_train_loss.data)\n",
    "\n",
    "        num_iter += 1\n",
    "\n",
    "    epoch_end_time = time.time()\n",
    "    per_epoch_ptime = epoch_end_time - epoch_start_time\n",
    "\n",
    "\n",
    "    print('[%d/%d] - ptime: %.2f, loss_d: %.3f, loss_g: %.3f' % ((epoch + 1), train_epoch, per_epoch_ptime, torch.mean(torch.FloatTensor(D_losses)),\n",
    "                                                              torch.mean(torch.FloatTensor(G_losses))))\n",
    "    p = 'MNIST_DCGAN_results/Random_results/MNIST_DCGAN_' + str(epoch + 1) + '.png'\n",
    "    fixed_p = 'MNIST_DCGAN_results/Fixed_results/MNIST_DCGAN_' + str(epoch + 1) + '.png'\n",
    "    show_result((epoch+1), save=True, path=p, isFix=False)\n",
    "    show_result((epoch+1), save=True, path=fixed_p, isFix=True)\n",
    "    train_hist['D_losses'].append(torch.mean(torch.FloatTensor(D_losses)))\n",
    "    train_hist['G_losses'].append(torch.mean(torch.FloatTensor(G_losses)))\n",
    "    train_hist['per_epoch_ptimes'].append(per_epoch_ptime)\n",
    "\n",
    "end_time = time.time()\n",
    "total_ptime = end_time - start_time\n",
    "train_hist['total_ptime'].append(total_ptime)\n",
    "\n",
    "print(\"Avg per epoch ptime: %.2f, total %d epochs ptime: %.2f\" % (torch.mean(torch.FloatTensor(train_hist['per_epoch_ptimes'])), train_epoch, total_ptime))\n",
    "print(\"Training finish!... save training results\")\n",
    "torch.save(G.state_dict(), \"MNIST_DCGAN_results/generator_param.pkl\")\n",
    "torch.save(D.state_dict(), \"MNIST_DCGAN_results/discriminator_param.pkl\")\n",
    "with open('MNIST_DCGAN_results/train_hist.pkl', 'wb') as f:\n",
    "    pickle.dump(train_hist, f)\n",
    "    \n",
    "    \n",
    "\n",
    "show_train_hist(train_hist, save=True, path='MNIST_DCGAN_results/MNIST_DCGAN_train_hist.png')\n",
    "\n",
    "images = []\n",
    "for e in range(train_epoch):\n",
    "    img_name = 'MNIST_DCGAN_results/Fixed_results/MNIST_DCGAN_' + str(e + 1) + '.png'\n",
    "    images.append(imageio.imread(img_name))\n",
    "imageio.mimsave('MNIST_DCGAN_results/generation_animation.gif', images, fps=5)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Few Shot Generation using the saved model\n",
    "\n",
    "### When saving a model for inference, it is only necessary to save the trained model’s learned parameters. Saving the model’s state_dict with the torch.save() function will give you the most flexibility for restoring the model later, which is why it is the recommended method for saving models.\n",
    "\n",
    "- Loading saved model and making predictions on nine\n",
    "\n",
    "To save models, we used:\n",
    "\n",
    "- torch.save(G.state_dict(), \"MNIST_DCGAN_results/generator_param.pkl\")\n",
    "- torch.save(D.state_dict(), \"MNIST_DCGAN_results/discriminator_param.pkl\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "metadata": {},
   "outputs": [],
   "source": [
    "def show_result(num_epoch, show = False, save = False, path = 'result.png', isFix=False):\n",
    "    z_ = torch.randn((5*5, 100)).view(-1, 100, 1, 1)\n",
    "    z_ = Variable(z_.cuda(), volatile=True)\n",
    "\n",
    "    g_model.eval()\n",
    "    if isFix:\n",
    "        test_images = g_model(fixed_z_)\n",
    "    else:\n",
    "        test_images = g_model(z_)\n",
    "    g_model.train()\n",
    "\n",
    "    size_figure_grid = 5\n",
    "    fig, ax = plt.subplots(size_figure_grid, size_figure_grid, figsize=(5, 5))\n",
    "    for i, j in itertools.product(range(size_figure_grid), range(size_figure_grid)):\n",
    "        ax[i, j].get_xaxis().set_visible(False)\n",
    "        ax[i, j].get_yaxis().set_visible(False)\n",
    "\n",
    "    for k in range(5*5):\n",
    "        i = k // 5\n",
    "        j = k % 5\n",
    "        ax[i, j].cla()\n",
    "        ax[i, j].imshow(test_images[k, 0].cpu().data.numpy(), cmap='gray')\n",
    "\n",
    "    label = 'Epoch {0}'.format(num_epoch)\n",
    "    fig.text(0.5, 0.04, label, ha='center')\n",
    "    plt.savefig(path)\n",
    "\n",
    "    if show:\n",
    "        plt.show()\n",
    "    else:\n",
    "        plt.close()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "discriminator(\n",
       "  (conv1): Conv2d(1, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2_bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv3_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv4_bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5): Conv2d(1024, 1, kernel_size=(4, 4), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_model = generator(128)\n",
    "g_model.load_state_dict(torch.load(\"MNIST_DCGAN_results/generator_param.pkl\"))\n",
    "g_model.eval()\n",
    "\n",
    "d_model = discriminator(128)\n",
    "d_model.load_state_dict(torch.load(\"MNIST_DCGAN_results/discriminator_param.pkl\"))\n",
    "d_model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "discriminator(\n",
       "  (conv1): Conv2d(1, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv2_bn): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv3): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv3_bn): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv4): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))\n",
       "  (conv4_bn): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)\n",
       "  (conv5): Conv2d(1024, 1, kernel_size=(4, 4), stride=(1, 1))\n",
       ")"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "g_model.cuda()\n",
    "d_model.cuda()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 60,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Print generator state_dict\n",
    "# print(\"Generator's state_dict:\")\n",
    "# for param_tensor in g_model.state_dict():\n",
    "#     print(param_tensor, \"\\t\", g_model.state_dict()[param_tensor].size())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 61,
   "metadata": {},
   "outputs": [],
   "source": [
    "# print(\"Discriminator's state_dict:\")\n",
    "# for param_tensor in d_model.state_dict():\n",
    "#     print(param_tensor, \"\\t\", d_model.state_dict()[param_tensor].size())  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 62,
   "metadata": {},
   "outputs": [],
   "source": [
    "# initializing parameters\n",
    "batch_size = 1\n",
    "lr = 0.0001\n",
    "train_epoch = 100\n",
    "img_size = 64\n",
    "\n",
    "# Binary Cross Entropy loss\n",
    "BCE_loss = nn.BCELoss()\n",
    "\n",
    "# Adam optimizer\n",
    "G_optimizer = optim.Adam(g_model.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "D_optimizer = optim.Adam(d_model.parameters(), lr=lr, betas=(0.5, 0.999))\n",
    "\n",
    "# training related parameters\n",
    "train_hist = {}\n",
    "train_hist['D_losses'] = []\n",
    "train_hist['G_losses'] = []\n",
    "train_hist['per_epoch_ptimes'] = []\n",
    "train_hist['total_ptime'] = []"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 63,
   "metadata": {},
   "outputs": [],
   "source": [
    "# results save folder\n",
    "if not os.path.isdir('NINE_results'):\n",
    "    os.mkdir('NINE_results')\n",
    "if not os.path.isdir('NINE_results/Random_results'):\n",
    "    os.mkdir('NINE_results/Random_results')\n",
    "if not os.path.isdir('NINE_results/Fixed_results'):\n",
    "    os.mkdir('NINE_results/Fixed_results')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 64,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Getting some images of nine ready for inference:\n",
    "\n",
    "transform = transforms.Compose([\n",
    "        transforms.ColorJitter(),\n",
    "        transforms.Scale(img_size),\n",
    "        transforms.ToTensor(),\n",
    "        transforms.Normalize((0.5,), (0.5,))\n",
    "])\n",
    "\n",
    "dataset = datasets.FashionMNIST(root='nine_data', train=True, download=True, transform=transform)\n",
    "\n",
    "idx = dataset.targets == 9\n",
    "\n",
    "dataset.data = dataset.data[idx]\n",
    "dataset.targets = dataset.targets[idx]\n",
    "\n",
    "train_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 65,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.utils.data.dataloader._DataLoaderIter'>\n",
      "images shape on batch size = torch.Size([1, 1, 64, 64])\n",
      "labels shape on batch size = torch.Size([1])\n"
     ]
    }
   ],
   "source": [
    "train_iter = iter(train_loader)\n",
    "print(type(train_iter))\n",
    "\n",
    "images, labels = train_iter.next()\n",
    "\n",
    "print('images shape on batch size = {}'.format(images.size()))\n",
    "print('labels shape on batch size = {}'.format(labels.size()))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 66,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Clipping input data to the valid range for imshow with RGB data ([0..1] for floats or [0..255] for integers).\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "Text(0.5, 1.0, '[9]')"
      ]
     },
     "execution_count": 66,
     "metadata": {},
     "output_type": "execute_result"
    },
    {
     "data": {
      "image/png": "iVBORw0KGgoAAAANSUhEUgAAAPsAAAEICAYAAACZA4KlAAAABHNCSVQICAgIfAhkiAAAAAlwSFlzAAALEgAACxIB0t1+/AAAADl0RVh0U29mdHdhcmUAbWF0cGxvdGxpYiB2ZXJzaW9uIDMuMC4zLCBodHRwOi8vbWF0cGxvdGxpYi5vcmcvnQurowAADPBJREFUeJzt3d1vVVUax/FnAZJKGFObgtUpaAhgQiqoQ0xVIhMumIY4vsUYg8RxHEfGxL+A8W4yF3M5CYnJhPHCWBNfowxBIUjUQaJVFBGxIBJfiqAxtjZgxFb2XAjLZy1Oz5zSfXra/r6fq2d37barwZ9rrbP3XjsURWEApr5pje4AgPFB2AERhB0QQdgBEYQdEEHYARGEHRBB2MWFEIoQwskQwt9rPP9PIYQTZ75vYb37h/IQdpiZLSuK4q9nD0IIvw8h7D8T6t0hhCVn24qi+HdRFLMb002MBWFHIoSwyMy6zewvZtZsZv8xs80hhBkN7RjGjLAj9zsz+29RFLuKohg2s3+Y2a/NbGVju4WxIuyoJGR1MLOOBvUFJSHsyO0ws5UhhN+GEGaa2QYzm2lmsxrbLYwVYUeiKIpeM/uDmW00s2Nm1mpmB8ysr5H9wtgFHnHVFkIozGxRURSHR2hvtp+DvvzM/whq+j5MPIzsOEcI4TchhOkhhDlm9i8z2+yDjsmJsKOSf5rZgJkdNLN+M/tzY7uDMhB2nDKzPSGEv539QlEUK4qi+FVRFC1FUawviuLk2bYQwh9DCANnvu90A/qL88SaHRDByA6IIOyAiHG93/nM5RoAdVQURaj0dUZ2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZABGEHRBB2QARhB0QQdkAEYQdEEHZAxIxGdwAT2yWXXBLrBQsWJG3t7e2xbmtrq+nnbd++PTk+ePDgGHqH0WBkB0QQdkAEYQdEsGaHLVu2LDnu7OyM9dKlS2O9ePHi5Dy/Tm9tbU3aZsz45T+tadN+GVOWL1+enPfYY4/Fevfu3Unb0NBQxf5eeumlybHv47Zt2yp+DxjZARmEHRARiqIYv18Wwvj9MiSuv/765Pjqq6+OdUdHR9I2f/78WPup+syZM5PzBgYGKtZmZrNmzapYf/PNN8l5X3755YhtP/zwQ6yHh4cr9snMbO7cubF++eWXk7YtW7bEur+/3xQURREqfZ2RHRBB2AERhB0QwaW3KezWW2+N9YYNG5I2v+696KKLkjZ/2ez06dOx/vrrr5Pz9u3bF+t33303aWtpaalY57fcrlmzJtb55Tt/yW6kOpf/LbNnz471o48+OuL3KWBkB0QQdkAE0/gp5vbbb4/16tWrY+0vT5mZnThxItZ9fX1J26FDh2L9zjvvVKzNzHp6emJd6yXca665Jjleu3ZtrG+44YakzV8CzPs/km+//TY57u3tren7FDCyAyIIOyCCsAMiWLNPcv7ympnZI488Emt/yStf8+7fvz/WBw4cSNqef/75WL/00kul9POs9957Lzn2l/n8bbVmZk1NTbE+3zW7//xBHSM7IIKwAyKYxk9CV111Vay7urqSNr8JpH8SLZ+qv/DCC7F++umnk7ZPPvmklH7W4oorroi177tZevebv5PPPw2XH+d/59GjR8vo5pTAyA6IIOyACKbxk8BDDz2UHN98882xXrhwYdLmHxLxd8b5T9/N0r3fvvrqq1L6WYsLLrggOb7pppti7TfUMEsfavHT+MHBweQ8/wn83r17S+nnVMTIDogg7IAIwg6IYM0+Qd1zzz2x9k+GmaUbROabNPoNJbZu3VqxNhvfdbrnn8QzS/+Wau+L80/p5U+y+afx3n777bF2ccpiZAdEEHZABNP4BspfY3T//ffH2m/k4DdxMEun7q+//nrStnPnzlh3d3eX0s+xuvHGG2OdT+NrfdWzn8bv2LEjadu4ceMYeqeDkR0QQdgBEYQdEMGafZxdeeWVsb7vvvuSNn/sbxXNn/Lye7Rv3749aXvqqadK6GW57rjjjlivWLEiafN7xfv3uZmZff/997H+/PPPY52/z+3kyZOl9HOqY2QHRBB2QATT+Dq7/PLLk+MHH3ww1kuXLk3a/H5s/lXGR44cSc7bvHlzrCfitN0sfbrNXzr0++KZpa+BPn78eNL25ptvxtrfAfj++++X1k8ljOyACMIOiGAaXwfz5s2L9cMPP5y0+Wm835Ah56fx+WuX/DR+ovIbUfhpfHNzc3Ken8bnb4n1Vxomyt2AkxkjOyCCsAMiCDsggjV7Hdx2222xXrJkSdLm16iffvpp0uY3iHz22WdjvWnTpuS8oaGhMrpZqvxVzP6uOX+XXP56Jv/ZRH5n3BNPPFFmF+UxsgMiCDsggml8CVauXJkc33LLLbH2rzcySx/28A93mKV3jD3++OOxbuS0/cILL4z1vffem7T5y2v53YD+rat+Gp+/nsn/zf7tsWZmP/3003n0GCNhZAdEEHZABGEHRLBmL8GCBQuS48WLF8fav3bYLL1FNn8vmV+nN3JDhhBCrB944IFY+88izNLPI/JNMUfiL7WZme3atSvWH3300Wi6iVFiZAdEEHZABNP4Evj94vJjf8ecWbqfnJ/Cmpl9/PHHdejdz/I96v1SI3/ts7/r79prrx3xvFmzZtX0u/3Sxd8laGb26quv1vQzMHaM7IAIwg6IYBpfAr93XH48bdrI/z/98ccfS+/LokWLYt3Z2RnrfAtn/+l5e3t70uY3mPBT9aamppr74e8UrHbX4HfffVfzz8TYMLIDIgg7IIKwAyJYs5dgcHAwOR4YGIh1fgedvxS3bt26pM1fsvOvZc5fi+TX1PnlMH9Xm1+L5+ty/1lCvvGl779/BVO+Zvd/W34Zzl9i9K9bzvfAx/hhZAdEEHZABNP4EuT7nfvXGPlNHMzM2traYn3nnXcmbXfddVfFn59P4/20O5+Cj9SWn+f7nPffXxL0y458SeIvMebTeD/99/vOHT582NAYjOyACMIOiCDsgAjW7CV48cUXk2N/2Slfl/s95fMn4vylLb/m9ZfCzNI1dn4pyx8fOnQo1r29vSP+jHwvd385r6urK9b+9luzdJ2ev8PNr817enpi/cEHHxgag5EdEEHYARFM4+tg27Ztsc43tli1alWsq91d56fx+aWxffv2xTrf/MEfHz16tPZOO3v27Kn49XyfOX8ZMb+056fx/tXLaBxGdkAEYQdEMI2vs2eeeSY59ps3VHuwxD9c4x+KMUs/cT916lQp/RyJvwMwf5VV/gm8x11zEw8jOyCCsAMiCDsggjX7OHvrrbca3YVR8fvG569lrrZvvP/MIX/lExqDkR0QQdgBEUzjUZW/S87XZuldc36fObN0A4+hoaE69Q6jwcgOiCDsgAjCDohgzY6qqu0v74/zTTH9xhmYGBjZARGEHRDBNB5V+el5rXvUm5n19fXVt2MYNUZ2QARhB0QwjUdV+TbTmLwY2QERhB0QQdgBEazZUVX+6ilMXozsgAjCDohgGo+q8tdXYfJiZAdEEHZABGEHRLBmR1UtLS2N7gJKwsgOiCDsgAim8TjH9OnTY+1fI43JjZEdEEHYARFM43GOuXPnxrqpqamBPUGZGNkBEYQdEEHYARGs2VGVf/0TJjf+JQERhB0QwTQe5zh27FisBwcHY11tSs90f+LjXwgQQdgBEYQdEMGaHVVV2zeeNfzkwr8IIIKwAyKYxqOqI0eOxLqvry9pq7anPPvNTzyM7IAIwg6IIOyACNbsqKramr29vT3W+caUrNknHkZ2QARhB0QwjUdVr7zySqyvu+66pG3VqlWxXr58edLW0dER69dee61OvcNoMLIDIgg7IIJpPGr25JNPJsdtbW2x7uzsTNr8lP/48eOxfu655+rUO/w/jOyACMIOiCDsgAjW7KjZZ599lhz39PTEuqurK2m77LLLYr1mzZpY+/W7mdkbb7xRZhdRBSM7IIKwAyKYxuO87dy5M9arV69O2lpbW2Pt76bzU3qzdI+7Dz/8sOwuwmFkB0QQdkAEYQdEsGbHefPvhNuyZUvS5p+IW7hwYazzp+OGh4djvWnTpqTtiy++KKWf+BkjOyCCsAMiQlEU4/fLQhi/X4aGmjdvXqzXrl0b6/nz5yfnNTc3x3rr1q1JW3d3d516N7UVRREqfZ2RHRBB2AERTONRd3PmzIn1unXrkra777471nv37k3a1q9fX9+OTVFM4wFxhB0QQdgBEazZ0VAXX3xxrPv7+xvYk6mDNTsgjrADIpjGA1MM03hAHGEHRBB2QARhB0QQdkAEYQdEjOulNwCNw8gOiCDsgAjCDogg7IAIwg6IIOyACMIOiCDsgAjCDogg7IAIwg6IIOyACMIOiCDsgAjCDogg7IAIwg6IIOyACMIOiCDsgAjCDogg7IAIwg6I+B/6GUwLyDEdjQAAAABJRU5ErkJggg==\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import torchvision\n",
    "\n",
    "grid = torchvision.utils.make_grid(images)\n",
    "\n",
    "plt.imshow(grid.numpy().transpose((1, 2, 0)))\n",
    "plt.axis('off')\n",
    "plt.title(labels.numpy())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 67,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "<class 'torch.utils.data.dataloader.DataLoader'>\n",
      "tensor([9])\n",
      "tensor([9])\n",
      "tensor([9])\n",
      "tensor([9])\n",
      "tensor([9])\n"
     ]
    }
   ],
   "source": [
    "# see if train_loader is giving only nines\n",
    "print(type(train_loader))\n",
    "\n",
    "count = 0\n",
    "\n",
    "for x_, y in train_loader:\n",
    "    count+= 1\n",
    "    # print(x_)\n",
    "    print(y)\n",
    "    if count == 5:\n",
    "        break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 68,
   "metadata": {},
   "outputs": [],
   "source": [
    "# create data to be trained upon: around 200 images\n",
    "\n",
    "train_custom_data = []\n",
    "\n",
    "for i in range(1, 52):\n",
    "    count = 0\n",
    "    for item in train_loader:\n",
    "        train_custom_data.append(item)\n",
    "        count += 1\n",
    "        if count == 4:\n",
    "            break"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 69,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "204\n"
     ]
    }
   ],
   "source": [
    "print(len(train_custom_data))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 70,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Fine tuning starts!\n",
      "epoch number 0\n",
      "count in epoch: 0 is: 204\n",
      "[1/100] - ptime: 3.69, loss_d: 0.457, loss_g: 9.983\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/home/jupyter/.local/lib/python3.5/site-packages/ipykernel_launcher.py:3: UserWarning: volatile was removed and now has no effect. Use `with torch.no_grad():` instead.\n",
      "  This is separate from the ipykernel package so we can avoid doing imports until\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "epoch number 1\n",
      "count in epoch: 1 is: 204\n",
      "[2/100] - ptime: 3.74, loss_d: 0.523, loss_g: 10.132\n",
      "epoch number 2\n",
      "count in epoch: 2 is: 204\n",
      "[3/100] - ptime: 3.72, loss_d: 0.217, loss_g: 9.160\n",
      "epoch number 3\n",
      "count in epoch: 3 is: 204\n",
      "[4/100] - ptime: 3.71, loss_d: 0.215, loss_g: 7.979\n",
      "epoch number 4\n",
      "count in epoch: 4 is: 204\n",
      "[5/100] - ptime: 3.72, loss_d: 0.581, loss_g: 6.350\n",
      "epoch number 5\n",
      "count in epoch: 5 is: 204\n",
      "[6/100] - ptime: 3.80, loss_d: 0.510, loss_g: 4.767\n",
      "epoch number 6\n",
      "count in epoch: 6 is: 204\n",
      "[7/100] - ptime: 3.79, loss_d: 0.627, loss_g: 2.917\n",
      "epoch number 7\n",
      "count in epoch: 7 is: 204\n",
      "[8/100] - ptime: 3.75, loss_d: 0.693, loss_g: 2.266\n",
      "epoch number 8\n",
      "count in epoch: 8 is: 204\n",
      "[9/100] - ptime: 3.76, loss_d: 0.710, loss_g: 1.988\n",
      "epoch number 9\n",
      "count in epoch: 9 is: 204\n",
      "[10/100] - ptime: 3.78, loss_d: 0.703, loss_g: 1.884\n",
      "epoch number 10\n",
      "count in epoch: 10 is: 204\n",
      "[11/100] - ptime: 3.79, loss_d: 0.706, loss_g: 1.826\n",
      "epoch number 11\n",
      "count in epoch: 11 is: 204\n",
      "[12/100] - ptime: 3.78, loss_d: 0.704, loss_g: 1.728\n",
      "epoch number 12\n",
      "count in epoch: 12 is: 204\n",
      "[13/100] - ptime: 3.80, loss_d: 0.704, loss_g: 1.715\n",
      "epoch number 13\n",
      "count in epoch: 13 is: 204\n",
      "[14/100] - ptime: 3.77, loss_d: 0.705, loss_g: 1.685\n",
      "epoch number 14\n",
      "count in epoch: 14 is: 204\n",
      "[15/100] - ptime: 3.72, loss_d: 0.695, loss_g: 1.662\n",
      "epoch number 15\n",
      "count in epoch: 15 is: 204\n",
      "[16/100] - ptime: 3.74, loss_d: 0.698, loss_g: 1.641\n",
      "epoch number 16\n",
      "count in epoch: 16 is: 204\n",
      "[17/100] - ptime: 3.74, loss_d: 0.695, loss_g: 1.639\n",
      "epoch number 17\n",
      "count in epoch: 17 is: 204\n",
      "[18/100] - ptime: 3.73, loss_d: 0.690, loss_g: 1.633\n",
      "epoch number 18\n",
      "count in epoch: 18 is: 204\n",
      "[19/100] - ptime: 3.77, loss_d: 0.688, loss_g: 1.625\n",
      "epoch number 19\n",
      "count in epoch: 19 is: 204\n",
      "[20/100] - ptime: 3.74, loss_d: 0.685, loss_g: 1.613\n",
      "epoch number 20\n",
      "count in epoch: 20 is: 204\n",
      "[21/100] - ptime: 3.75, loss_d: 0.682, loss_g: 1.611\n",
      "epoch number 21\n",
      "count in epoch: 21 is: 204\n",
      "[22/100] - ptime: 3.78, loss_d: 0.686, loss_g: 1.608\n",
      "epoch number 22\n",
      "count in epoch: 22 is: 204\n",
      "[23/100] - ptime: 3.74, loss_d: 0.684, loss_g: 1.604\n",
      "epoch number 23\n",
      "count in epoch: 23 is: 204\n",
      "[24/100] - ptime: 3.74, loss_d: 0.676, loss_g: 1.602\n",
      "epoch number 24\n",
      "count in epoch: 24 is: 204\n",
      "[25/100] - ptime: 3.75, loss_d: 0.674, loss_g: 1.601\n",
      "epoch number 25\n",
      "count in epoch: 25 is: 204\n",
      "[26/100] - ptime: 3.80, loss_d: 0.673, loss_g: 1.601\n",
      "epoch number 26\n",
      "count in epoch: 26 is: 204\n",
      "[27/100] - ptime: 3.67, loss_d: 0.671, loss_g: 1.597\n",
      "epoch number 27\n",
      "count in epoch: 27 is: 204\n",
      "[28/100] - ptime: 3.78, loss_d: 0.669, loss_g: 1.595\n",
      "epoch number 28\n",
      "count in epoch: 28 is: 204\n",
      "[29/100] - ptime: 3.79, loss_d: 0.670, loss_g: 1.595\n",
      "epoch number 29\n",
      "count in epoch: 29 is: 204\n",
      "[30/100] - ptime: 3.78, loss_d: 0.669, loss_g: 1.592\n",
      "epoch number 30\n",
      "count in epoch: 30 is: 204\n",
      "[31/100] - ptime: 3.82, loss_d: 0.663, loss_g: 1.594\n",
      "epoch number 31\n",
      "count in epoch: 31 is: 204\n",
      "[32/100] - ptime: 3.76, loss_d: 0.662, loss_g: 1.592\n",
      "epoch number 32\n",
      "count in epoch: 32 is: 204\n",
      "[33/100] - ptime: 3.78, loss_d: 0.667, loss_g: 1.590\n",
      "epoch number 33\n",
      "count in epoch: 33 is: 204\n",
      "[34/100] - ptime: 3.78, loss_d: 0.659, loss_g: 1.591\n",
      "epoch number 34\n",
      "count in epoch: 34 is: 204\n",
      "[35/100] - ptime: 3.77, loss_d: 0.658, loss_g: 1.591\n",
      "epoch number 35\n",
      "count in epoch: 35 is: 204\n",
      "[36/100] - ptime: 3.72, loss_d: 0.657, loss_g: 1.592\n",
      "epoch number 36\n",
      "count in epoch: 36 is: 204\n",
      "[37/100] - ptime: 3.80, loss_d: 0.654, loss_g: 1.593\n",
      "epoch number 37\n",
      "count in epoch: 37 is: 204\n",
      "[38/100] - ptime: 3.82, loss_d: 0.654, loss_g: 1.593\n",
      "epoch number 38\n",
      "count in epoch: 38 is: 204\n",
      "[39/100] - ptime: 3.78, loss_d: 0.841, loss_g: 1.681\n",
      "epoch number 39\n",
      "count in epoch: 39 is: 204\n",
      "[40/100] - ptime: 3.81, loss_d: 0.656, loss_g: 1.586\n",
      "epoch number 40\n",
      "count in epoch: 40 is: 204\n",
      "[41/100] - ptime: 3.81, loss_d: 0.650, loss_g: 1.589\n",
      "epoch number 41\n",
      "count in epoch: 41 is: 204\n",
      "[42/100] - ptime: 3.81, loss_d: 0.647, loss_g: 1.592\n",
      "epoch number 42\n",
      "count in epoch: 42 is: 204\n",
      "[43/100] - ptime: 3.82, loss_d: 0.646, loss_g: 1.593\n",
      "epoch number 43\n",
      "count in epoch: 43 is: 204\n",
      "[44/100] - ptime: 3.77, loss_d: 0.645, loss_g: 1.594\n",
      "epoch number 44\n",
      "count in epoch: 44 is: 204\n",
      "[45/100] - ptime: 3.74, loss_d: 0.644, loss_g: 1.594\n",
      "epoch number 45\n",
      "count in epoch: 45 is: 204\n",
      "[46/100] - ptime: 3.72, loss_d: 0.645, loss_g: 1.593\n",
      "epoch number 46\n",
      "count in epoch: 46 is: 204\n",
      "[47/100] - ptime: 3.82, loss_d: 0.644, loss_g: 1.594\n",
      "epoch number 47\n",
      "count in epoch: 47 is: 204\n",
      "[48/100] - ptime: 3.88, loss_d: 0.643, loss_g: 1.595\n",
      "epoch number 48\n",
      "count in epoch: 48 is: 204\n",
      "[49/100] - ptime: 3.79, loss_d: 0.642, loss_g: 1.596\n",
      "epoch number 49\n",
      "count in epoch: 49 is: 204\n",
      "[50/100] - ptime: 3.80, loss_d: 0.642, loss_g: 1.596\n",
      "epoch number 50\n",
      "count in epoch: 50 is: 204\n",
      "[51/100] - ptime: 3.80, loss_d: 0.641, loss_g: 1.596\n",
      "epoch number 51\n",
      "count in epoch: 51 is: 204\n",
      "[52/100] - ptime: 3.78, loss_d: 0.643, loss_g: 1.596\n",
      "epoch number 52\n",
      "count in epoch: 52 is: 204\n",
      "[53/100] - ptime: 3.83, loss_d: 0.912, loss_g: 2.480\n",
      "epoch number 53\n",
      "count in epoch: 53 is: 204\n",
      "[54/100] - ptime: 3.77, loss_d: 0.674, loss_g: 1.603\n",
      "epoch number 54\n",
      "count in epoch: 54 is: 204\n",
      "[55/100] - ptime: 3.80, loss_d: 0.656, loss_g: 1.597\n",
      "epoch number 55\n",
      "count in epoch: 55 is: 204\n",
      "[56/100] - ptime: 3.76, loss_d: 0.650, loss_g: 1.598\n",
      "epoch number 56\n",
      "count in epoch: 56 is: 204\n",
      "[57/100] - ptime: 3.80, loss_d: 0.647, loss_g: 1.598\n",
      "epoch number 57\n",
      "count in epoch: 57 is: 204\n",
      "[58/100] - ptime: 3.74, loss_d: 0.645, loss_g: 1.599\n",
      "epoch number 58\n",
      "count in epoch: 58 is: 204\n",
      "[59/100] - ptime: 3.78, loss_d: 0.643, loss_g: 1.598\n",
      "epoch number 59\n",
      "count in epoch: 59 is: 204\n",
      "[60/100] - ptime: 3.77, loss_d: 0.642, loss_g: 1.600\n",
      "epoch number 60\n",
      "count in epoch: 60 is: 204\n",
      "[61/100] - ptime: 3.78, loss_d: 0.642, loss_g: 1.599\n",
      "epoch number 61\n",
      "count in epoch: 61 is: 204\n",
      "[62/100] - ptime: 3.74, loss_d: 0.641, loss_g: 1.599\n",
      "epoch number 62\n",
      "count in epoch: 62 is: 204\n",
      "[63/100] - ptime: 3.78, loss_d: 0.641, loss_g: 1.599\n",
      "epoch number 63\n",
      "count in epoch: 63 is: 204\n",
      "[64/100] - ptime: 3.81, loss_d: 0.640, loss_g: 1.599\n",
      "epoch number 64\n",
      "count in epoch: 64 is: 204\n",
      "[65/100] - ptime: 3.80, loss_d: 0.641, loss_g: 1.598\n",
      "epoch number 65\n",
      "count in epoch: 65 is: 204\n",
      "[66/100] - ptime: 3.82, loss_d: 0.639, loss_g: 1.599\n",
      "epoch number 66\n",
      "count in epoch: 66 is: 204\n",
      "[67/100] - ptime: 3.81, loss_d: 0.640, loss_g: 1.598\n",
      "epoch number 67\n",
      "count in epoch: 67 is: 204\n",
      "[68/100] - ptime: 3.80, loss_d: 0.639, loss_g: 1.599\n",
      "epoch number 68\n",
      "count in epoch: 68 is: 204\n",
      "[69/100] - ptime: 3.82, loss_d: 0.644, loss_g: 1.596\n",
      "epoch number 69\n",
      "count in epoch: 69 is: 204\n",
      "[70/100] - ptime: 3.82, loss_d: 0.650, loss_g: 1.592\n",
      "epoch number 70\n",
      "count in epoch: 70 is: 204\n",
      "[71/100] - ptime: 3.79, loss_d: 0.639, loss_g: 1.596\n",
      "epoch number 71\n",
      "count in epoch: 71 is: 204\n",
      "[72/100] - ptime: 3.78, loss_d: 0.638, loss_g: 1.597\n",
      "epoch number 72\n",
      "count in epoch: 72 is: 204\n",
      "[73/100] - ptime: 3.82, loss_d: 0.637, loss_g: 1.599\n",
      "epoch number 73\n",
      "count in epoch: 73 is: 204\n",
      "[74/100] - ptime: 3.85, loss_d: 0.637, loss_g: 1.599\n",
      "epoch number 74\n",
      "count in epoch: 74 is: 204\n",
      "[75/100] - ptime: 3.80, loss_d: 0.637, loss_g: 1.599\n",
      "epoch number 75\n",
      "count in epoch: 75 is: 204\n",
      "[76/100] - ptime: 3.81, loss_d: 0.636, loss_g: 1.599\n",
      "epoch number 76\n",
      "count in epoch: 76 is: 204\n",
      "[77/100] - ptime: 3.72, loss_d: 0.637, loss_g: 1.599\n",
      "epoch number 77\n",
      "count in epoch: 77 is: 204\n",
      "[78/100] - ptime: 3.72, loss_d: 0.636, loss_g: 1.598\n",
      "epoch number 78\n",
      "count in epoch: 78 is: 204\n",
      "[79/100] - ptime: 3.77, loss_d: 0.641, loss_g: 1.596\n",
      "epoch number 79\n",
      "count in epoch: 79 is: 204\n",
      "[80/100] - ptime: 3.76, loss_d: 0.635, loss_g: 1.599\n",
      "epoch number 80\n",
      "count in epoch: 80 is: 204\n",
      "[81/100] - ptime: 3.74, loss_d: 0.635, loss_g: 1.599\n",
      "epoch number 81\n",
      "count in epoch: 81 is: 204\n",
      "[82/100] - ptime: 3.69, loss_d: 0.637, loss_g: 1.600\n",
      "epoch number 82\n",
      "count in epoch: 82 is: 204\n",
      "[83/100] - ptime: 3.79, loss_d: 0.636, loss_g: 1.602\n",
      "epoch number 83\n",
      "count in epoch: 83 is: 204\n",
      "[84/100] - ptime: 3.74, loss_d: 4.699, loss_g: 1.526\n",
      "epoch number 84\n",
      "count in epoch: 84 is: 204\n",
      "[85/100] - ptime: 3.71, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 85\n",
      "count in epoch: 85 is: 204\n",
      "[86/100] - ptime: 3.74, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 86\n",
      "count in epoch: 86 is: 204\n",
      "[87/100] - ptime: 3.71, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 87\n",
      "count in epoch: 87 is: 204\n",
      "[88/100] - ptime: 3.78, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 88\n",
      "count in epoch: 88 is: 204\n",
      "[89/100] - ptime: 3.73, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 89\n",
      "count in epoch: 89 is: 204\n",
      "[90/100] - ptime: 3.72, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 90\n",
      "count in epoch: 90 is: 204\n",
      "[91/100] - ptime: 3.73, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 91\n",
      "count in epoch: 91 is: 204\n",
      "[92/100] - ptime: 3.76, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 92\n",
      "count in epoch: 92 is: 204\n",
      "[93/100] - ptime: 3.70, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 93\n",
      "count in epoch: 93 is: 204\n",
      "[94/100] - ptime: 3.75, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 94\n",
      "count in epoch: 94 is: 204\n",
      "[95/100] - ptime: 3.70, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 95\n",
      "count in epoch: 95 is: 204\n",
      "[96/100] - ptime: 3.76, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 96\n",
      "count in epoch: 96 is: 204\n",
      "[97/100] - ptime: 3.75, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 97\n",
      "count in epoch: 97 is: 204\n",
      "[98/100] - ptime: 3.72, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 98\n",
      "count in epoch: 98 is: 204\n",
      "[99/100] - ptime: 3.74, loss_d: 27.631, loss_g: 0.000\n",
      "epoch number 99\n",
      "count in epoch: 99 is: 204\n",
      "[100/100] - ptime: 3.72, loss_d: 27.631, loss_g: 0.000\n"
     ]
    }
   ],
   "source": [
    "num_iter = 0\n",
    "\n",
    "print('Fine tuning starts!')\n",
    "\n",
    "start_time = time.time()\n",
    "\n",
    "for epoch in range(train_epoch):\n",
    "    \n",
    "    print(\"epoch number\", epoch)\n",
    "    D_losses = []\n",
    "    G_losses = []\n",
    "    epoch_start_time = time.time()\n",
    "    \n",
    "    # Few shot, remember?\n",
    "    count = 0\n",
    "    \n",
    "    for x_, _ in train_custom_data:\n",
    "        \n",
    "        # train discriminator D\n",
    "        d_model.zero_grad()\n",
    "\n",
    "        mini_batch = x_.size()[0]\n",
    "\n",
    "        y_real_ = torch.ones(mini_batch)\n",
    "        y_fake_ = torch.zeros(mini_batch)\n",
    "\n",
    "        x_, y_real_, y_fake_ = Variable(x_.cuda()), Variable(y_real_.cuda()), Variable(y_fake_.cuda())\n",
    "        D_result = d_model(x_).squeeze()\n",
    "        D_real_loss = BCE_loss(D_result, y_real_)\n",
    "\n",
    "        z_ = torch.randn((mini_batch, 100)).view(-1, 100, 1, 1)\n",
    "        z_ = Variable(z_.cuda())\n",
    "        G_result = g_model(z_)\n",
    "\n",
    "        D_result = d_model(G_result).squeeze()\n",
    "        D_fake_loss = BCE_loss(D_result, y_fake_)\n",
    "        D_fake_score = D_result.data.mean()\n",
    "\n",
    "        D_train_loss = D_real_loss + D_fake_loss\n",
    "\n",
    "        D_train_loss.backward()\n",
    "        D_optimizer.step()\n",
    "\n",
    "        # D_losses.append(D_train_loss.data[0])\n",
    "        D_losses.append(D_train_loss.data)\n",
    "\n",
    "        # train generator G\n",
    "        g_model.zero_grad()\n",
    "\n",
    "        z_ = torch.randn((mini_batch, 100)).view(-1, 100, 1, 1)\n",
    "        z_ = Variable(z_.cuda())\n",
    "\n",
    "        G_result = g_model(z_)\n",
    "        D_result = d_model(G_result).squeeze()\n",
    "        G_train_loss = BCE_loss(D_result, y_real_)\n",
    "        G_train_loss.backward()\n",
    "        G_optimizer.step()\n",
    "\n",
    "        G_losses.append(G_train_loss.data)\n",
    "\n",
    "        num_iter += 1\n",
    "        count += 1\n",
    "\n",
    "    print(\"count in epoch: \" + str(epoch) + \" is: \" + str(count))    \n",
    "    epoch_end_time = time.time()\n",
    "    per_epoch_ptime = epoch_end_time - epoch_start_time\n",
    "\n",
    "\n",
    "    print('[%d/%d] - ptime: %.2f, loss_d: %.3f, loss_g: %.3f' % ((epoch + 1), train_epoch, per_epoch_ptime, torch.mean(torch.FloatTensor(D_losses)),\n",
    "                                                              torch.mean(torch.FloatTensor(G_losses))))\n",
    "    p = 'NINE_results/Random_results/NINE_DCGAN_' + str(epoch + 1) + '.png'\n",
    "    fixed_p = 'NINE_results/Fixed_results/NINE_DCGAN_' + str(epoch + 1) + '.png'\n",
    "    show_result((epoch+1), save=True, path=p, isFix=False)\n",
    "    show_result((epoch+1), save=True, path=fixed_p, isFix=True)\n",
    "    train_hist['D_losses'].append(torch.mean(torch.FloatTensor(D_losses)))\n",
    "    train_hist['G_losses'].append(torch.mean(torch.FloatTensor(G_losses)))\n",
    "    train_hist['per_epoch_ptimes'].append(per_epoch_ptime)\n",
    "\n",
    "end_time = time.time()\n",
    "total_ptime = end_time - start_time\n",
    "train_hist['total_ptime'].append(total_ptime)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
